import onnx
import torch
import caffe2.python.onnx.backend as backend
from chooseModel import choose
import numpy as np

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

# torch.load('filename.pth').to(device)

# model = torch.load(
#     '/home/ly/Projects/pytorch_my/checkpoints_color_regress_adam_resnet50_cv_blur_5x5/checkpoint_115.pth', map_location=device)
model = choose('resnet50', 1, False)
checkpoint = torch.load('/home/ly/Projects/pytorch_my/checkpoints_color_regress_adam_resnet50_cv_blur_5x5/checkpoint_115.pth')
model.load_state_dict(checkpoint['state_dict'])

model.eval()
model.cuda()
# batch_size = 1  # 批处理大小
input_shape = (1, 3, 224, 224)  # 输入数据

dummy_input = torch.randn(input_shape, device=device)
model(dummy_input)

torch.onnx.export(model, dummy_input, "resnet50_cv_blur_5x5.onnx", verbose=False)
with torch.no_grad():
    output_1 = model(dummy_input)
print(output_1)

# Model check after conversion
model = onnx.load("resnet50_cv_blur_5x5.onnx")
rep = backend.prepare(model, device="CUDA:0")
output_2 = rep.run(dummy_input.cpu().numpy().astype(np.float32))
print(output_2)
try:
    onnx.checker.check_model(model)
    print('ONNX check passed successfully.')
except onnx.onnx_cpp2py_export.checker.ValidationError as exc:
    sys.exit('ONNX check failed with error: ' + str(exc))
